Summary

• serialize the adjoint traced-key paths into the forward simulation’s HDF5 metadata (__tidy3d_traced_field_keys__) so cache hashes differ when the tracked set changes
• move the Tracer, FieldMap, and TracerKeys helpers into tidy3d/components/autograd and reuse them from both the serializer and the web pipeline
• add a regression test confirming simulations with identical statics but different tracer sets now hash differently

Alternatives Considered

• keep relying on the existing sidecar TracerKeys file: this leaves the cache blind to tracer changes, because the hash sees identical HDF5 exports
• expose traced keys as a new model field: that would surface runtime autograd state in the public schema even though users neither set nor persist it
• store the paths in model.attrs: those attrs are user-controlled metadata and cached JSON already assumes they are non-functional

Rationale

Writing the traced-key summary to an HDF5 attribute keeps the metadata with the exported simulation, feeds into the hashing logic, and avoids making autograd bookkeeping part of the user-facing schema.

Greptile Overview

Updated On: 2025-10-08 08:26:25 UTC

Summary

The solution serializes traced field keys directly into the HDF5 metadata using a new attribute __tidy3d_traced_field_keys__. This ensures that simulations tracking different autograd parameters generate different cache hashes, preventing incorrect cache hits that could lead to wrong gradient computations in optimization workflows.

The implementation includes a significant refactoring that moves the Tracer, FieldMap, and TracerKeys helper classes from tidy3d/web/api/autograd/utils.py into a new shared module tidy3d/components/autograd/field_map.py. This centralization enables code reuse between the web pipeline and the serialization components while maintaining consistency in traced field metadata handling.

A new regression test test_sim_hash_changes_with_traced_keys() validates that when the traced field set changes (by converting a structure to static), both the field map and the simulation hash change accordingly, ensuring proper cache invalidation.

Important Files Changed

Confidence score: 4/5

• This PR addresses a critical correctness issue in autograd caching with a well-reasoned solution
• Score reflects solid implementation with proper testing, though the core change touches critical caching logic
• Pay close attention to tidy3d/components/base.py where the HDF5 metadata serialization is implemented

Sequence Diagram

sequenceDiagram
    participant User
    participant Simulation
    participant BaseModel as "Tidy3dBaseModel"
    participant FieldMap as "FieldMap/TracerKeys"
    participant HDF5 as "HDF5 File"
    participant Cache as "Cache System"

    User->>Simulation: "Create simulation with traced fields"
    Simulation->>BaseModel: "_strip_traced_fields()"
    BaseModel-->>Simulation: "AutogradFieldMap"
    
    User->>Simulation: "Export to HDF5"
    Simulation->>BaseModel: "to_hdf5()"
    BaseModel->>BaseModel: "_serialized_traced_field_keys()"
    BaseModel->>FieldMap: "TracerKeys.from_field_mapping()"
    FieldMap-->>BaseModel: "TracerKeys instance"
    BaseModel->>FieldMap: "tracer_keys.json()"
    FieldMap-->>BaseModel: "JSON serialized keys"
    BaseModel->>HDF5: "Store traced keys in attributes"
    BaseModel->>HDF5: "Write simulation data"
    
    User->>Cache: "Check cache hash"
    Cache->>BaseModel: "_hash_self()"
    BaseModel->>BaseModel: "to_hdf5() (in memory)"
    Note over BaseModel,HDF5: "HDF5 now includes traced keys in metadata"
    BaseModel-->>Cache: "Hash includes traced keys"
    
    Note over User,Cache: "Different traced sets now produce different hashes"